from typing import TypedDict

from langgraph.constants import START, END
from langgraph.graph import StateGraph


class InputState(TypedDict):
    user_input: str


class OutputState(TypedDict):
    graph_output: str


class OverallState(TypedDict):
    foo: str
    user_input: str
    graph_output: str


class PrivateState(TypedDict):
    bar: str



if __name__ == '__main__':
    def node_1(state: InputState) -> OverallState:
        # Write to OverallState
        return {"foo": state["user_input"] + " name"}


    def node_2(state: OverallState) -> PrivateState:
        # Read from OverallState, write to PrivateState
        return {"bar": state["foo"] + " is"}


    def node_3(state: PrivateState) -> OutputState:
        # Read from PrivateState, write to OutputState
        return {"graph_output": state["bar"] + " Lance"}


    builder = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState)
    builder.add_node("node_1", node_1)
    builder.add_node("node_2", node_2)
    builder.add_node("node_3", node_3)
    builder.add_edge(START, "node_1")
    builder.add_edge("node_1", "node_2")
    builder.add_edge("node_2", "node_3")
    builder.add_edge("node_3", END)

    graph = builder.compile()
    print(graph.invoke({"user_input": "My"}))
    # {'graph_output': 'My name is Lance'}
